{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "03_siamese_triplet_mnist.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [
        "JpWMG8kF1Ezv",
        "e5ib15cj1Ezy",
        "OIjXt_931E0A",
        "GfA7R5dF1E0F",
        "sNEzScHM1E0F",
        "5Z1fGqzm1E0M",
        "EpVpPbBZ1E0N",
        "l7nclgG91E0Q",
        "xOzgmRO91E0X",
        "Ncftmsls1E0b",
        "azGVIKYV1E0e",
        "yCRCt5Fj1E0q",
        "779Hh05A1E0u",
        "PyCDVrG81E0w",
        "AsMAV5Kt1E0z",
        "LEkIWNDq1E05",
        "X13jYJfk1E1R"
      ],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mlelarge/dataflowr/blob/master/Notebooks/03_siamese/03_siamese_triplet_mnist_colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "metadata": {
        "id": "Ei8ToiWt1EzS",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Siamese networks"
      ]
    },
    {
      "metadata": {
        "id": "kd38ctpK1EzT",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Colab preparation"
      ]
    },
    {
      "metadata": {
        "id": "jEKc30JO1EzU",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# to install pytorch on colab\n",
        "from os import path\n",
        "from wheel.pep425tags import get_abbr_impl, get_impl_ver, get_abi_tag\n",
        "platform = '{}{}-{}'.format(get_abbr_impl(), get_impl_ver(), get_abi_tag())\n",
        "\n",
        "accelerator = 'cu80' if path.exists('/opt/bin/nvidia-smi') else 'cpu'\n",
        "\n",
        "!pip install -q http://download.pytorch.org/whl/{accelerator}/torch-1.0.0-{platform}-linux_x86_64.whl torchvision"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "sx9e_pXlCuti",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "%load_ext autoreload\n",
        "%autoreload 2"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "aU__26HN1EzZ",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import random\n",
        "\n",
        "%matplotlib inline\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import torch\n",
        "from torch.optim import lr_scheduler\n",
        "import torch.optim as optim\n",
        "from torch.autograd import Variable\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision\n",
        "from torchvision.datasets import MNIST\n",
        "from torchvision import transforms\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "UMMut8UVCutt"
      },
      "cell_type": "markdown",
      "source": [
        "# 1. Setup and initializations\n",
        "We'll go through learning feature embeddings using different loss functions on MNIST dataset. This is just for visualization purposes, thus we'll be using 2-dimensional embeddings which isn't the best choice in practice.\n",
        "\n",
        "For every experiment the same embedding network is used (`32 conv 5x5 -> ReLU -> MaxPool 2x2 -> 64 conv 5x5 -> ReLU -> MaxPool 2x2 -> Fully Connected 256 -> ReLU -> Fully Connected 256 -> ReLU -> Fully Connected 2`) with the same hyperparameters."
      ]
    },
    {
      "metadata": {
        "id": "tvM4mzUm1Ezc",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class ExperimentParams():\n",
        "    def __init__(self):\n",
        "        self.num_classes = 10\n",
        "        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "        self.batch_size = 256\n",
        "        self.lr = 1e-2\n",
        "        self.num_epochs = 10\n",
        "        self.num_workers = 4\n",
        "        self.data_dir = '/home/docker_user/'\n",
        "        \n",
        "\n",
        "args = ExperimentParams()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "BcmGBqXeCutw"
      },
      "cell_type": "markdown",
      "source": [
        "## 1.1 Prepare dataset\n",
        "We'll be working on MNIST dataset"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "AQSHPH9P0BNx",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\n",
        "mean, std = 0.1307, 0.3081\n",
        "\n",
        "train_dataset = MNIST(f'{args.data_dir}/data/MNIST', train=True, download=True,\n",
        "                             transform=transforms.Compose([\n",
        "                                 transforms.ToTensor(),\n",
        "                                 transforms.Normalize((mean,), (std,))\n",
        "                             ]))\n",
        "test_dataset = MNIST(f'{args.data_dir}/data/MNIST', train=False, download=True,\n",
        "                            transform=transforms.Compose([\n",
        "                                transforms.ToTensor(),\n",
        "                                transforms.Normalize((mean,), (std,))\n",
        "                            ]))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "TcZTFRnjCut3"
      },
      "cell_type": "markdown",
      "source": [
        "## 1.2 Common setup"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "Dz2xh66UCut5",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\n",
        "mnist_classes = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']\n",
        "colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728',\n",
        "              '#9467bd', '#8c564b', '#e377c2', '#7f7f7f',\n",
        "              '#bcbd22', '#17becf']\n",
        "\n",
        "def plot_embeddings(embeddings, targets, title='',xlim=None, ylim=None):\n",
        "    plt.figure(figsize=(10,10))\n",
        "    for i in range(10):\n",
        "        inds = np.where(targets==i)[0]\n",
        "        plt.scatter(embeddings[inds,0], embeddings[inds,1], alpha=0.5, color=colors[i])\n",
        "    if xlim:\n",
        "        plt.xlim(xlim[0], xlim[1])\n",
        "    if ylim:\n",
        "        plt.ylim(ylim[0], ylim[1])\n",
        "    plt.legend(mnist_classes)\n",
        "    plt.title(title)\n",
        "\n",
        "def extract_embeddings(dataloader, model, args):\n",
        "    with torch.no_grad():\n",
        "        model.eval()\n",
        "        embeddings = np.zeros((len(dataloader.dataset), 2))\n",
        "        labels = np.zeros(len(dataloader.dataset))\n",
        "        k = 0\n",
        "        for images, target in dataloader:\n",
        "            images = images.to(args.device)\n",
        "            embeddings[k:k+len(images)] = model.get_embedding(images).data.cpu().numpy()\n",
        "            labels[k:k+len(images)] = target.numpy()\n",
        "            k += len(images)\n",
        "    return embeddings, labels\n",
        "\n",
        "\n",
        "def get_raw_images(dataloader,mean=0.1307, std=0.3081):\n",
        "\n",
        "    raw_images = np.zeros((len(dataloader.dataset), 1, 28, 28))\n",
        "    k = 0\n",
        "    for input, target in dataloader:\n",
        "        raw_images[k:k+len(input)] = (input*std + mean).data.cpu().numpy()\n",
        "        k += len(input)\n",
        "\n",
        "    return raw_images\n",
        "\n",
        "\n",
        "def show(img, title=None):\n",
        "    # img is a torch.Tensor     \n",
        "    npimg = img.numpy()\n",
        "    plt.imshow(np.transpose(npimg, (1,2,0)), interpolation='nearest')\n",
        "    plt.axis('off')\n",
        "    if title is not None:\n",
        "        plt.title(title)\n",
        "    plt.pause(0.001)  # pause a bit so that plots are updated"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "75moY8AyCut_"
      },
      "cell_type": "markdown",
      "source": [
        "# 2. Baseline: Classification with softmax\n",
        "We'll train the model for classification and use outputs of penultimate layer as embeddings. "
      ]
    },
    {
      "metadata": {
        "id": "zJCmiKwf1Ezl",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "We will define our base embedding architecture which will serve as common backbone for our experiments"
      ]
    },
    {
      "metadata": {
        "id": "Ew1qgU_G1Ezl",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 2.1 Architecture"
      ]
    },
    {
      "metadata": {
        "id": "49bn26xC1Ezm",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "#### Exercise\n",
        "\n",
        "Complete the missing blocks in the definition of the following `EmbeddingNet` architecture: (`32 conv 5x5 -> ReLU -> MaxPool 2x2 -> 64 conv 5x5 -> ReLU -> MaxPool 2x2 -> Fully Connected 256 -> ReLU -> Fully Connected 256 -> ReLU -> Fully Connected 2`)"
      ]
    },
    {
      "metadata": {
        "id": "C1LjKH_21Ezn",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class EmbeddingNet(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(EmbeddingNet, self).__init__()\n",
        "#         self.conv1 = nn.Conv2d(1, ...)\n",
        "#         self.conv2 = ...\n",
        "#         self.fc1 = ...\n",
        "#         self.fc2 = ...\n",
        "#         self.fc3 = ...\n",
        "\n",
        "    def forward(self, x, debug=False):\n",
        "        x1 = F.max_pool2d(F.relu(self.conv1(x)), kernel_size=2, stride=2)\n",
        "#         output = ...                  \n",
        "        if debug == True:\n",
        "            print(f'input: {x.size()}')\n",
        "            print(f'x1: {x1.size()}')\n",
        "                      \n",
        "        return output\n",
        "\n",
        "    def get_embedding(self, x):\n",
        "        return self.forward(x)\n",
        "    \n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "lOYKYEm41Ezq",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "If you want to better check the sizes of the hidden states and do debugging, you can add a `debug` variable in the `forward` function just like above"
      ]
    },
    {
      "metadata": {
        "id": "D1PZbEVB1Ezr",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "input = torch.zeros(1, 1, 28, 28)\n",
        "net = EmbeddingNet()\n",
        "net(input,debug=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "uGMMtw_F1Ezv",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Now let's define a classification net that will add fully connected layer on top of `EmbeddingNet`"
      ]
    },
    {
      "metadata": {
        "id": "JpWMG8kF1Ezv",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "#### Exercice\n",
        "\n",
        "Fill in the missing spots in the `forward` pass:"
      ]
    },
    {
      "metadata": {
        "id": "7ChpUMXh1Ezw",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class ClassificationNet(nn.Module):\n",
        "    def __init__(self, embedding_net, num_classes):\n",
        "        super(ClassificationNet, self).__init__()\n",
        "        self.embedding_net = embedding_net\n",
        "        self.prelu = nn.PReLU()\n",
        "        self.fc = nn.Linear(2, num_classes)\n",
        "\n",
        "    def forward(self, x, debug=False):\n",
        "        embedding = None\n",
        "        output = self.fc(self.prelu(embedding))\n",
        "        \n",
        "        if debug == True:\n",
        "            print(f'input: {x.size()}')\n",
        "            print(f'embedding: {embedding.size()}')\n",
        "            print(f'output: {output.size()}')\n",
        "            \n",
        "        return output\n",
        "    \n",
        "    def get_embedding(self, x):\n",
        "        return self.prelu(None)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "e5ib15cj1Ezy",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 2.2 Training"
      ]
    },
    {
      "metadata": {
        "id": "KA8szJpj1Ezy",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Set up data loaders\n",
        "\n",
        "kwargs = {'num_workers': args.num_workers, 'pin_memory': True} \n",
        "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, **kwargs)\n",
        "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, **kwargs)\n",
        "\n",
        "embedding_net = EmbeddingNet()\n",
        "model = ClassificationNet(embedding_net, num_classes=args.num_classes)\n",
        "loss_fn = torch.nn.CrossEntropyLoss()\n",
        "model.to(args.device)\n",
        "loss_fn.to(args.device)\n",
        "\n",
        "optimizer = optim.Adam(model.parameters(), lr=args.lr)\n",
        "scheduler = lr_scheduler.StepLR(optimizer, 8, gamma=0.1, last_epoch=-1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "i_NBcq3w1Ez3",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "train_embeddings_baseline, train_labels_baseline = extract_embeddings(train_loader, model, args)\n",
        "plot_embeddings(train_embeddings_baseline, train_labels_baseline, 'Train embeddings before training')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "eOLUpe-U1Ez6",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def train_classif_epoch(train_loader, model, loss_fn, optimizer, args, log_interval=50):\n",
        "    model.train()\n",
        "    losses = []\n",
        "    total_loss, total_corrects, num_samples = 0, 0, 0\n",
        "    corrects = 0    \n",
        "    for batch_idx, (data, target) in enumerate(train_loader):\n",
        "        num_samples += data.size(0)\n",
        "        \n",
        "        data, target = data.to(args.device), target.to(args.device)\n",
        "        \n",
        "        optimizer.zero_grad()\n",
        "        outputs = model(data)\n",
        "\n",
        "        loss = loss_fn(outputs, target)\n",
        "        losses.append(loss.data.item())\n",
        "\n",
        "        _,preds = torch.max(outputs.data,1)\n",
        "        corrects += torch.sum(preds == target.data).cpu()\n",
        "\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        \n",
        "        if batch_idx % log_interval == 0:\n",
        "            print('Train: [{}/{} ({:.0f}%)]\\tLoss: {:.6f} \\tAccuracy: {}'.format(\n",
        "                batch_idx * len(data[0]), len(train_loader.dataset),\n",
        "                100. * batch_idx / len(train_loader), np.mean(losses), float(total_corrects)/num_samples))           \n",
        "            \n",
        "            total_loss += np.sum(losses)\n",
        "            total_corrects += corrects\n",
        "            losses, corrects = [], 0\n",
        "\n",
        "    return total_loss/(batch_idx + 1), total_corrects/num_samples\n",
        "\n",
        "def test_classif_epoch(test_loader, model, loss_fn, args, log_interval=50):\n",
        "    with torch.no_grad():\n",
        "        model.eval()\n",
        "        losses, corrects = [], 0\n",
        "        num_samples = 0\n",
        "    \n",
        "        for batch_idx, (data, target) in enumerate(test_loader):\n",
        "\n",
        "            num_samples += data.size(0)\n",
        "            data, target = data.to(args.device), target.to(args.device)\n",
        "\n",
        "            outputs = model(data)\n",
        "\n",
        "            loss = loss_fn(outputs, target)\n",
        "            losses.append(loss.data.item())\n",
        "\n",
        "            _,preds = torch.max(outputs.data,1)\n",
        "            corrects += torch.sum(preds == target.data).cpu()\n",
        "\n",
        "        return np.sum(losses)/(batch_idx + 1), corrects/num_samples"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "yfGNHVAz1Ez9",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "start_epoch = 0\n",
        "\n",
        "for epoch in range(0, start_epoch):\n",
        "    scheduler.step()\n",
        "\n",
        "for epoch in range(start_epoch, args.num_epochs):\n",
        "    scheduler.step()\n",
        "\n",
        "    train_loss, train_accuracy = train_classif_epoch(train_loader, model, loss_fn, optimizer, args)\n",
        "\n",
        "    message = 'Epoch: {}/{}. Train set: Average loss: {:.4f} Average accuracy: {:.4f}'.format(\n",
        "        epoch + 1, args.num_epochs, train_loss, train_accuracy)\n",
        "    \n",
        "    val_loss, val_accuracy = test_classif_epoch(test_loader, model, loss_fn, args)\n",
        "    \n",
        "    message += '\\nEpoch: {}/{}. Validation set: Average loss: {:.4f}  Average accuracy: {:.4f}'.format(epoch + 1, args.num_epochs,\n",
        "                                                                             val_loss, val_accuracy)\n",
        "    print(message)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "OIjXt_931E0A",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 2.3 Visualizations"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "q9o-6-ytCuuT",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "train_embeddings_baseline, train_labels_baseline = extract_embeddings(train_loader, model, args)\n",
        "plot_embeddings(train_embeddings_baseline, train_labels_baseline, 'Train embeddings classification')\n",
        "test_embeddings_baseline, test_labels_baseline = extract_embeddings(test_loader, model, args)\n",
        "plot_embeddings(test_embeddings_baseline, test_labels_baseline, 'Test embeddings classification')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "fDdsCIiHCuub"
      },
      "cell_type": "markdown",
      "source": [
        "While the embeddings look separable (which is what we trained them for), they don't have good metric properties. They might not be the best choice as a descriptor for new classes."
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "j9FhHE-tCuuc"
      },
      "cell_type": "markdown",
      "source": [
        "# 3. Siamese network\n",
        "Now we'll train a siamese network that takes a pair of images and trains the embeddings so that the distance between them is minimized if their from the same class or greater than some margin value if they represent different classes.\n",
        "We'll minimize a contrastive loss function*:\n",
        "$$L_{contrastive}(x_0, x_1, y) = \\frac{1}{2} y \\lVert f(x_0)-f(x_1)\\rVert_2^2 + \\frac{1}{2}(1-y)\\{max(0, m-\\lVert f(x_0)-f(x_1)\\rVert_2)\\}^2$$\n",
        "\n",
        "*Raia Hadsell, Sumit Chopra, Yann LeCun, [Dimensionality reduction by learning an invariant mapping](http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf), CVPR 2006*"
      ]
    },
    {
      "metadata": {
        "id": "GfA7R5dF1E0F",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 3.1 Architecture\n",
        "We will first define the siamese architecture on top of our `EmbeddingNet`"
      ]
    },
    {
      "metadata": {
        "id": "sNEzScHM1E0F",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "#### Exercise\n",
        "\n",
        "Fill in the forward part of `SiameseNet`"
      ]
    },
    {
      "metadata": {
        "id": "Hj7Jzzwz1E0G",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class SiameseNet(nn.Module):\n",
        "    def __init__(self, embedding_net):\n",
        "        super(SiameseNet, self).__init__()\n",
        "        self.embedding_net = embedding_net\n",
        "\n",
        "    def forward(self, x1, x2):\n",
        "        # fill in the missing 2 lines :)\n",
        "        \n",
        "        return output1, output2\n",
        "\n",
        "    def get_embedding(self, x):\n",
        "        return self.embedding_net(x)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "935Hwn0R1E0J",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 3.2 Data loader\n",
        "We will also need to adapt our data loader to fetch pairs of images "
      ]
    },
    {
      "metadata": {
        "id": "XAFAu0We1E0K",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from torch.utils.data import Dataset\n",
        "from torch.utils.data.sampler import BatchSampler\n",
        "from PIL import Image\n",
        "\n",
        "class SiameseMNIST(Dataset):\n",
        "    \"\"\"\n",
        "    train mode: For each sample creates randomly a positive or a negative pair\n",
        "    test mode: Creates fixed pairs for testing\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, mnist_dataset):\n",
        "        self.mnist_dataset = mnist_dataset\n",
        "\n",
        "        self.train = self.mnist_dataset.train\n",
        "        self.transform = self.mnist_dataset.transform\n",
        "\n",
        "        if self.train:\n",
        "            self.train_labels = self.mnist_dataset.train_labels\n",
        "            self.train_data = self.mnist_dataset.train_data\n",
        "            self.labels_set = set(self.train_labels.numpy())\n",
        "            self.label_to_indices = {label: np.where(self.train_labels.numpy() == label)[0]\n",
        "                                     for label in self.labels_set}\n",
        "        else:\n",
        "            # generate fixed pairs for testing\n",
        "            self.test_labels = self.mnist_dataset.test_labels\n",
        "            self.test_data = self.mnist_dataset.test_data\n",
        "            self.labels_set = set(self.test_labels.numpy())\n",
        "            '''\n",
        "            create a dictionary with an entry key for each label and the value an array storing\n",
        "            the indices of the images having the respective label\n",
        "            '''\n",
        "            self.label_to_indices = {label: np.where(self.test_labels.numpy() == label)[0]\n",
        "                                     for label in self.labels_set}\n",
        "\n",
        "            random_state = np.random.RandomState(42)\n",
        "            # itereate through test_data and randomly select samples with the same label\n",
        "            positive_pairs = [[i,\n",
        "                               random_state.choice(self.label_to_indices[self.test_labels[i].item()]),\n",
        "                               1]\n",
        "                              for i in range(0, len(self.test_data), 2)]\n",
        "\n",
        "            # itereate through test_data, create a list of all labels different from current one and then\n",
        "            # randomly select samples with having one of these labels\n",
        "            negative_pairs = [[i,\n",
        "                               random_state.choice(self.label_to_indices[\n",
        "                                                       np.random.choice(\n",
        "                                                           list(self.labels_set - set([self.test_labels[i].item()]))\n",
        "                                                       )\n",
        "                                                   ]),\n",
        "                               0]\n",
        "                              for i in range(1, len(self.test_data), 2)]\n",
        "            # format: [index1, index2, label(0/1)]\n",
        "            self.test_pairs = positive_pairs + negative_pairs\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        \n",
        "        # at train time pairs of samples are fetched randomly on the fly\n",
        "        if self.train:\n",
        "            # select random label,i.e. similar (1) or non-similar (0) images\n",
        "            target = np.random.randint(0, 2)\n",
        "            img1, label1 = self.train_data[index], self.train_labels[index].item()\n",
        "            if target == 1:\n",
        "                # select an image with the same label as img1\n",
        "                siamese_index = index\n",
        "                while siamese_index == index:\n",
        "                    siamese_index = np.random.choice(self.label_to_indices[label1])\n",
        "            else:\n",
        "                # eliminate label1 from the set of possible labels to select\n",
        "                siamese_label = np.random.choice(list(self.labels_set - set([label1])))\n",
        "                # randomly select an image having a label from this subset\n",
        "                siamese_index = np.random.choice(self.label_to_indices[siamese_label])\n",
        "            img2 = self.train_data[siamese_index]\n",
        "        else:\n",
        "            img1 = self.test_data[self.test_pairs[index][0]]\n",
        "            img2 = self.test_data[self.test_pairs[index][1]]\n",
        "            target = self.test_pairs[index][2]\n",
        "\n",
        "        img1 = Image.fromarray(img1.numpy(), mode='L')\n",
        "        img2 = Image.fromarray(img2.numpy(), mode='L')\n",
        "        if self.transform is not None:\n",
        "            img1 = self.transform(img1)\n",
        "            img2 = self.transform(img2)\n",
        "        return (img1, img2), target\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.mnist_dataset)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "5Z1fGqzm1E0M",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 3.3 Loss function"
      ]
    },
    {
      "metadata": {
        "id": "C48BKMc91E0N",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "$$L_{contrastive}(x_0, x_1, y) = \\frac{1}{2} y \\lVert f(x_0)-f(x_1)\\rVert_2^2 + \\frac{1}{2}(1-y)\\{max(0, m-\\lVert f(x_0)-f(x_1)\\rVert_2)\\}^2$$"
      ]
    },
    {
      "metadata": {
        "id": "EpVpPbBZ1E0N",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "#### Exercise\n",
        "\n",
        "Fill in the missing parts of the `contrastive loss`"
      ]
    },
    {
      "metadata": {
        "id": "jiKkNs3e1E0O",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class ContrastiveLoss(nn.Module):\n",
        "    \"\"\"\n",
        "    Contrastive loss\n",
        "    Takes embeddings of two samples and a target label == 1 if samples are from the same class and label == 0 otherwise\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, margin):\n",
        "        super(ContrastiveLoss, self).__init__()\n",
        "        self.margin = margin\n",
        "        self.eps = 1e-9\n",
        "\n",
        "    def forward(self, output1, output2, target, size_average=True):\n",
        "        # compute squared distances between output2 and output1  \n",
        "        squared_distances = None\n",
        "        # add the second term from them loss. You can use ReLU for compressing the max formula\n",
        "        losses = 0.5 * (target.float() * squared_distances +\n",
        "                         None )\n",
        "        \n",
        "        return losses.mean() if size_average else losses.sum()\n",
        "        \n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "l7nclgG91E0Q",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 3.4 Training"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "-qpIq-TzCuue",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Set up data loaders\n",
        "siamese_train_dataset = SiameseMNIST(train_dataset) # Returns pairs of images and target same/different\n",
        "siamese_test_dataset = SiameseMNIST(test_dataset)\n",
        "\n",
        "args.batch_size = 128\n",
        "kwargs = {'num_workers': args.num_workers, 'pin_memory': True}\n",
        "siamese_train_loader = torch.utils.data.DataLoader(siamese_train_dataset, batch_size=args.batch_size, shuffle=True, **kwargs)\n",
        "siamese_test_loader = torch.utils.data.DataLoader(siamese_test_dataset, batch_size=args.batch_size, shuffle=False, **kwargs)\n",
        "\n",
        "margin = 1.\n",
        "embedding_net = EmbeddingNet()\n",
        "model = SiameseNet(embedding_net)\n",
        "loss_fn = ContrastiveLoss(margin)\n",
        "model.to(args.device)\n",
        "loss_fn.to(args.device)\n",
        "\n",
        "args.lr = 1e-3\n",
        "optimizer = optim.Adam(model.parameters(), lr=args.lr)\n",
        "scheduler = lr_scheduler.StepLR(optimizer, 8, gamma=0.1, last_epoch=-1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "fYJbkKAm1E0S",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def train_siamese_epoch(train_loader, model, loss_fn, optimizer, args, log_interval=100):\n",
        "    model.train()\n",
        "    losses = []\n",
        "    total_loss, num_samples =  0, 0\n",
        "  \n",
        "    for batch_idx, (data, target) in enumerate(train_loader):\n",
        "        num_samples += data[0].size(0)\n",
        "        \n",
        "        data = tuple(d.to(args.device) for d in data)\n",
        "        target = target.to(args.device)\n",
        "          \n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        outputs = model(data[0], data[1])\n",
        "        # alternatively: outputs = model(*data)\n",
        "        \n",
        "        loss = loss_fn(outputs[0], outputs[1], target)\n",
        "        # alternatively: loss = loss_fn(*outputs, target)\n",
        "        \n",
        "        losses.append(loss.data.item())\n",
        "\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        \n",
        "        if batch_idx % log_interval == 0:\n",
        "            print('Train: [{}/{} ({:.0f}%)]\\tLoss: {:.6f} '.format(\n",
        "                batch_idx * len(data[0]), len(train_loader.dataset),\n",
        "                100. * batch_idx / len(train_loader), np.mean(losses)))           \n",
        "            \n",
        "            total_loss += np.sum(losses)\n",
        "            losses = []\n",
        "            \n",
        "    return total_loss/(batch_idx + 1)\n",
        "\n",
        "def test_siamese_epoch(test_loader, model, loss_fn, args, log_interval=50):\n",
        "    with torch.no_grad():\n",
        "        model.eval()\n",
        "        losses = []\n",
        "        num_samples = 0\n",
        "    \n",
        "        for batch_idx, (data, target) in enumerate(test_loader):\n",
        "\n",
        "            num_samples += data[0].size(0)\n",
        "            data = tuple(d.to(args.device) for d in data)\n",
        "            target = target.to(args.device)\n",
        "            outputs = model(data[0], data[1])\n",
        "\n",
        "            loss = loss_fn(outputs[0], outputs[1], target)\n",
        "            losses.append(loss.data.item())\n",
        "    \n",
        "        return np.sum(losses)/(batch_idx + 1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "_IqiBATeCuuh",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "start_epoch = 0\n",
        "\n",
        "# needed for annealing learning rate in case of resuming of training\n",
        "for epoch in range(0, start_epoch):\n",
        "    scheduler.step()\n",
        "\n",
        "# main training loop\n",
        "for epoch in range(start_epoch, args.num_epochs):\n",
        "    scheduler.step()\n",
        "\n",
        "    # train stage\n",
        "    train_loss = train_siamese_epoch(siamese_train_loader, model, loss_fn, optimizer, args)\n",
        "    message = 'Epoch: {}/{}. Train set: Average loss: {:.4f}'.format(\n",
        "        epoch + 1, args.num_epochs, train_loss)\n",
        "    \n",
        "    # testing/validation stage    \n",
        "    test_loss = test_siamese_epoch(siamese_test_loader, model, loss_fn, args)\n",
        "    \n",
        "    message += '\\nEpoch: {}/{}. Validation set: Average loss: {:.4f}'.format(epoch + 1, args.num_epochs,\n",
        "                                                                             test_loss)\n",
        "    print(message)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "xOzgmRO91E0X",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 3.5 Visualizations"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "rBBBuPjiCuup",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "train_embeddings_cl, train_labels_cl = extract_embeddings(train_loader, model, args)\n",
        "plot_embeddings(train_embeddings_cl, train_labels_cl, title='Train embeddings (constrastive loss)')\n",
        "test_embeddings_cl, test_labels_cl = extract_embeddings(test_loader, model, args)\n",
        "plot_embeddings(test_embeddings_cl, test_labels_cl, title='Test embeddings (contrastive loss)')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "28NPJHe91E0b",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "In order to two compare vectors $x_1$ and $x_2$ we can use the `cosine similarity` \n",
        "\n",
        "$$\\text{similarity}=\\frac{x_1 \\cdot x_2}{\\max(\\Vert x_1 \\Vert _2 \\cdot \\Vert x_2 \\Vert_2, \\epsilon)}$$\n",
        "\n",
        "An alternative is the Euclidean distance."
      ]
    },
    {
      "metadata": {
        "id": "ObKaI7Oa1E0b",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "In order to save computation at query time we can pre-process our vectors and L2-normalize them. Now we can simply perform comparison by dot product"
      ]
    },
    {
      "metadata": {
        "id": "Ncftmsls1E0b",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "#### Exercise\n",
        "Perform L2-normalization on the embeddings using `numpy` "
      ]
    },
    {
      "metadata": {
        "id": "rigYIlwM1E0c",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# L2-normalize embeddings\n",
        "test_embeddings_norm = ...."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "azGVIKYV1E0e",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "#### Exercise\n",
        "Write now a function `most_sim` that computes all dot products between a query vector and the dataset, extracts the indices of the `topk` most similar vectors and put thme in a list of tuples ("
      ]
    },
    {
      "metadata": {
        "id": "HuvS4r6D1E0f",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def most_sim(x, emb, topk=6):\n",
        "       return None"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "hF5DQ3F71E0h",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "test_images_raw = get_raw_images(test_loader)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "TzxTj6jn1E0i",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def launch_query(test_embeddings_norm, test_images_raw, query_id=None):\n",
        "    query_id = random.randint(0, test_embeddings_norm.shape[0]) if query_id is None else query_id\n",
        "    query_vector = test_embeddings_norm[query_id,:]\n",
        "\n",
        "    print(f'query_id: {query_id} | query_embedding: {query_vector}')\n",
        "    knns = most_sim(query_vector, test_embeddings_norm)\n",
        "    knn_images = np.array([test_images_raw[x[0]] for x in knns ])\n",
        "    \n",
        "    title=['q: 1.0', f'1nn: {knns[1][1]:.3}', f'2nn: {knns[2][1]:.3}', \n",
        "           f'3nn: {knns[3][1]:.3}', f'4nn: {knns[4][1]:.3}', f'5nn: {knns[5][1]:.3}']\n",
        "    show(torchvision.utils.make_grid(torch.from_numpy(knn_images)), title=title)\n",
        "#     print(knns)\n",
        "    "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ydZMwGHt1E0k",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "for i in range(5):\n",
        "    launch_query(test_embeddings_norm, test_images_raw)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "MbKXy6yQCuuu"
      },
      "cell_type": "markdown",
      "source": [
        "# Triplet network\n",
        "We'll train a triplet network, that takes an anchor, positive (same class as anchor) and negative (different class than anchor) examples. The objective is to learn embeddings such that the anchor is closer to the positive example than it is to the negative example by some margin value.\n",
        "\n",
        "![alt text](images/anchor_negative_positive.png \"Source: FaceNet\")\n",
        "Source: [2] *Schroff, Florian, Dmitry Kalenichenko, and James Philbin. [Facenet: A unified embedding for face recognition and clustering.](https://arxiv.org/abs/1503.03832) CVPR 2015.*\n",
        "\n",
        "**Triplet loss**:   $L_{triplet}(x_a, x_p, x_n) = max(0, m +  \\lVert f(x_a)-f(x_p)\\rVert_2^2 - \\lVert f(x_a)-f(x_n)\\rVert_2^2$\\)"
      ]
    },
    {
      "metadata": {
        "id": "pBLyLplg1E0m",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 4.1 Architecture\n",
        "We will first define the triplet architecture on top of our `EmbeddingNet`\n",
        "\n",
        "#### Exercise\n",
        "\n",
        "Fill in the forward part of `TripleNet`"
      ]
    },
    {
      "metadata": {
        "id": "b2kAbre51E0o",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class TripletNet(nn.Module):\n",
        "    def __init__(self, embedding_net):\n",
        "        super(TripletNet, self).__init__()\n",
        "        self.embedding_net = embedding_net\n",
        "\n",
        "    def forward(self, x1, x2, x3):\n",
        "        # missing 3 lines here\n",
        "        \n",
        "        return output1, output2, output3\n",
        "\n",
        "    def get_embedding(self, x):\n",
        "        return self.embedding_net(x)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "yCRCt5Fj1E0q",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 4.2 Data loader\n",
        "We will also need to adapt our data loader to fetch triplets of images "
      ]
    },
    {
      "metadata": {
        "id": "MvAuRt6J1E0s",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from torch.utils.data import Dataset\n",
        "from torch.utils.data.sampler import BatchSampler\n",
        "from PIL import Image\n",
        "\n",
        "class TripletMNIST(Dataset):\n",
        "    \"\"\"\n",
        "    Train: For each sample (anchor) randomly chooses a positive and negative samples\n",
        "    Test: Creates fixed triplets for testing\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, mnist_dataset):\n",
        "        self.mnist_dataset = mnist_dataset\n",
        "        self.train = self.mnist_dataset.train\n",
        "        self.transform = self.mnist_dataset.transform\n",
        "\n",
        "        if self.train:\n",
        "            self.train_labels = self.mnist_dataset.train_labels\n",
        "            self.train_data = self.mnist_dataset.train_data\n",
        "            self.labels_set = set(self.train_labels.numpy())\n",
        "            self.label_to_indices = {label: np.where(self.train_labels.numpy() == label)[0]\n",
        "                                     for label in self.labels_set}\n",
        "\n",
        "        else:\n",
        "            self.test_labels = self.mnist_dataset.test_labels\n",
        "            self.test_data = self.mnist_dataset.test_data\n",
        "            # generate fixed triplets for testing\n",
        "            self.labels_set = set(self.test_labels.numpy())\n",
        "            self.label_to_indices = {label: np.where(self.test_labels.numpy() == label)[0]\n",
        "                                     for label in self.labels_set}\n",
        "\n",
        "            random_state = np.random.RandomState(29)\n",
        "\n",
        "            triplets = [[i,\n",
        "                         random_state.choice(self.label_to_indices[self.test_labels[i].item()]),\n",
        "                         random_state.choice(self.label_to_indices[\n",
        "                                                 np.random.choice(\n",
        "                                                     list(self.labels_set - set([self.test_labels[i].item()]))\n",
        "                                                 )\n",
        "                                             ])\n",
        "                         ]\n",
        "                        for i in range(len(self.test_data))]\n",
        "            self.test_triplets = triplets\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        if self.train:\n",
        "            img1, label1 = self.train_data[index], self.train_labels[index].item()\n",
        "            positive_index = index\n",
        "            while positive_index == index:\n",
        "                positive_index = np.random.choice(self.label_to_indices[label1])\n",
        "            negative_label = np.random.choice(list(self.labels_set - set([label1])))\n",
        "            negative_index = np.random.choice(self.label_to_indices[negative_label])\n",
        "            img2 = self.train_data[positive_index]\n",
        "            img3 = self.train_data[negative_index]\n",
        "        else:\n",
        "            img1 = self.test_data[self.test_triplets[index][0]]\n",
        "            img2 = self.test_data[self.test_triplets[index][1]]\n",
        "            img3 = self.test_data[self.test_triplets[index][2]]\n",
        "\n",
        "        img1 = Image.fromarray(img1.numpy(), mode='L')\n",
        "        img2 = Image.fromarray(img2.numpy(), mode='L')\n",
        "        img3 = Image.fromarray(img3.numpy(), mode='L')\n",
        "        if self.transform is not None:\n",
        "            img1 = self.transform(img1)\n",
        "            img2 = self.transform(img2)\n",
        "            img3 = self.transform(img3)\n",
        "        return (img1, img2, img3), []\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.mnist_dataset)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "779Hh05A1E0u",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 4.3 Loss function\n",
        "\n",
        "#### Exercise\n",
        "\n",
        "Fill in the missing parts of the `triplet loss`:\n",
        " $L_{triplet}(x_a, x_p, x_n) = max(0, m +  \\lVert f(x_a)-f(x_p)\\rVert_2^2 - \\lVert f(x_a)-f(x_n)\\rVert_2^2$\\)"
      ]
    },
    {
      "metadata": {
        "id": "kRaxvIFJ1E0u",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class TripletLoss(nn.Module):\n",
        "    \"\"\"\n",
        "    Triplet loss\n",
        "    Takes embeddings of an anchor sample, a positive sample and a negative sample\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, margin):\n",
        "        super(TripletLoss, self).__init__()\n",
        "        self.margin = margin\n",
        "\n",
        "    def forward(self, anchor, positive, negative, size_average=True):\n",
        "        distance_positive = None # fill in code\n",
        "        distance_negative = None  # fill in code\n",
        "        # you can again use ReLU instead of max         \n",
        "        losses = None # fill in code\n",
        "        return losses.mean() if size_average else losses.sum()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "PyCDVrG81E0w",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 4.4 Training"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "jv4DvFucCuuu",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "triplet_train_dataset = TripletMNIST(train_dataset) # Returns triplets of images\n",
        "triplet_test_dataset = TripletMNIST(test_dataset)\n",
        "\n",
        "args.batch_size = 128\n",
        "kwargs = {'num_workers': args.num_workers, 'pin_memory': True} \n",
        "triplet_train_loader = torch.utils.data.DataLoader(triplet_train_dataset, batch_size=args.batch_size, shuffle=True, **kwargs)\n",
        "triplet_test_loader = torch.utils.data.DataLoader(triplet_test_dataset, batch_size=args.batch_size, shuffle=False, **kwargs)\n",
        "\n",
        "margin = 1.\n",
        "embedding_net = EmbeddingNet()\n",
        "model = TripletNet(embedding_net)\n",
        "loss_fn = TripletLoss(margin)\n",
        "model.to(args.device)\n",
        "loss_fn.to(args.device)\n",
        "args.lr = 1e-3\n",
        "optimizer = optim.Adam(model.parameters(), lr=args.lr)\n",
        "scheduler = lr_scheduler.StepLR(optimizer, 8, gamma=0.1, last_epoch=-1)\n",
        "n_epochs = 5\n",
        "log_interval = 100"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "AsMAV5Kt1E0z",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "#### Exercice\n",
        "\n",
        "Code your own train/test sequences similarly to the previous examples.\n",
        "Watch out for some differences though."
      ]
    },
    {
      "metadata": {
        "id": "HIeuIH4y1E00",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def train_triplet_epoch(train_loader, model, loss_fn, optimizer, args, log_interval=100):\n",
        "    model.train()\n",
        "    losses = []\n",
        "    total_loss, num_samples =  0, 0\n",
        "  \n",
        "    # fill in code here\n",
        "    \n",
        "    return total_loss/(batch_idx + 1)\n",
        "\n",
        "def test_triplet_epoch(test_loader, model, loss_fn, args, log_interval=50):\n",
        "    losses = []\n",
        "    num_samples = 0\n",
        "    # fill in code here\n",
        "\n",
        "    return np.sum(losses)/(batch_idx + 1)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "6TfLZH1V1E02",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "start_epoch = 0\n",
        "\n",
        "# needed for annealing learning rate in case of resuming of training\n",
        "for epoch in range(0, start_epoch):\n",
        "    scheduler.step()\n",
        "\n",
        "# main training loop\n",
        "for epoch in range(start_epoch, args.num_epochs):\n",
        "    scheduler.step()\n",
        "\n",
        "    # train stage\n",
        "    train_loss = train_triplet_epoch(triplet_train_loader, model, loss_fn, optimizer, args)\n",
        "    message = 'Epoch: {}/{}. Train set: Average loss: {:.4f}'.format(\n",
        "        epoch + 1, args.num_epochs, train_loss)\n",
        "    \n",
        "    # testing/validation stage    \n",
        "    test_loss = test_triplet_epoch(triplet_test_loader, model, loss_fn, args)\n",
        "    \n",
        "    message += '\\nEpoch: {}/{}. Validation set: Average loss: {:.4f}'.format(epoch + 1, args.num_epochs,\n",
        "                                                                             test_loss)\n",
        "    print(message)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "LEkIWNDq1E05",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 4.5 Visualizations"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "ysh4Ry7ZCuu_",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "train_embeddings_tl, train_labels_tl = extract_embeddings(train_loader, model, args)\n",
        "plot_embeddings(train_embeddings_tl, train_labels_tl, title='Train triplet embeddings')\n",
        "test_embeddings_tl, test_labels_tl = extract_embeddings(test_loader, model, args)\n",
        "plot_embeddings(test_embeddings_tl, test_labels_tl, title='Val triplet embeddings')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "zw0y8bgB1E08",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# L2-normalize embeddings\n",
        "test_embeddings_tl_norm = test_embeddings_tl / np.linalg.norm(test_embeddings_tl, axis=-1, keepdims=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "YnNl2ffL1E09",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "test_images_raw = get_raw_images(test_loader)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "h9JOhZTY1E0_",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "for i in range(5):\n",
        "    launch_query(test_embeddings_tl_norm, test_images_raw)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "x7C9H1_nCuvJ"
      },
      "cell_type": "markdown",
      "source": [
        "# 5. Online pair selection - negative mining\n",
        "There are couple of problems with siamese (and triplet) networks.\n",
        "1. The number of possible pairs/triplets grows **quadratically/cubically** with the number of examples. It's infeasible to process them all\n",
        "2. We generate pairs/triplets randomly. As the training continues, more and more pairs/triplets are easy to deal with (their loss value is very small or even 0), preventing the network from training. We need to provide the network with **hard examples**.\n",
        "3. Each image that is fed to the network is used only for computation of contrastive/triplet loss for only one pair/triplet. The computation is somewhat wasted; once the embedding is computed, it could be reused for many pairs/triplets.\n",
        "\n",
        "To deal with that efficiently, we'll feed a network with standard mini-batches as we did for classification. The loss function will be responsible for selection of hard pairs and triplets within mini-batch. In these case, if we feed the network with 16 images per 10 classes, we can process up to $159*160/2 = 12720$ pairs compared to 80 pairs in previous implementation.\n"
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "k806qej9CuvL"
      },
      "cell_type": "markdown",
      "source": [
        "## 5.1 Online pair selection"
      ]
    },
    {
      "metadata": {
        "id": "7pT7ftyT1E1C",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class BalancedBatchSampler(BatchSampler):\n",
        "    \"\"\"\n",
        "    BatchSampler - from a MNIST-like dataset, samples num_classes and within these classes samples num_samples.\n",
        "    Returns batches of size num_classes * num_samples\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, dataset, num_classes, num_samples):\n",
        "        if dataset.train:\n",
        "            self.labels = dataset.train_labels\n",
        "        else:\n",
        "            self.labels = dataset.test_labels\n",
        "            \n",
        "        self.labels_set = list(set(self.labels.numpy()))\n",
        "        self.label_to_indices = {label: np.where(self.labels.numpy() == label)[0]\n",
        "                                 for label in self.labels_set}\n",
        "        for l in self.labels_set:\n",
        "            np.random.shuffle(self.label_to_indices[l])\n",
        "            \n",
        "        self.used_label_indices_count = {label: 0 for label in self.labels_set}\n",
        "        self.count = 0\n",
        "        self.num_classes = num_classes\n",
        "        self.num_samples = num_samples\n",
        "        self.dataset = dataset\n",
        "        self.batch_size = self.num_samples * self.num_classes\n",
        "\n",
        "    def __iter__(self):\n",
        "        self.count = 0\n",
        "        while self.count + self.batch_size < len(self.dataset):\n",
        "            classes = np.random.choice(self.labels_set, self.num_classes, replace=False)\n",
        "            indices = []\n",
        "            for class_ in classes:\n",
        "                indices.extend(self.label_to_indices[class_][\n",
        "                               self.used_label_indices_count[class_]:self.used_label_indices_count[\n",
        "                                                                         class_] + self.num_samples])\n",
        "                self.used_label_indices_count[class_] += self.num_samples\n",
        "                if self.used_label_indices_count[class_] + self.num_samples > len(self.label_to_indices[class_]):\n",
        "                    np.random.shuffle(self.label_to_indices[class_])\n",
        "                    self.used_label_indices_count[class_] = 0\n",
        "            yield indices\n",
        "            self.count += self.num_classes * self.num_samples\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.dataset) // self.batch_size"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "cQh1Of9M1E1E",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 5.2 Online pair selection"
      ]
    },
    {
      "metadata": {
        "id": "XE9dKb3K1E1H",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from itertools import combinations\n",
        "\n",
        "def pdist(vectors):\n",
        "    distance_matrix = -2 * vectors.mm(torch.t(vectors)) + vectors.pow(2).sum(dim=1).view(1, -1) + vectors.pow(2).sum(\n",
        "        dim=1).view(-1, 1)\n",
        "    return distance_matrix\n",
        "\n",
        "\n",
        "class HardNegativePairSelector():\n",
        "    \"\"\"\n",
        "    Creates all possible positive pairs. For negative pairs, pairs with smallest distance are taken into consideration,\n",
        "    matching the number of positive pairs.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, cpu=True):\n",
        "#         super(HardNegativePairSelector, self).__init__()\n",
        "        self.cpu = cpu\n",
        "\n",
        "    def get_pairs(self, embeddings, labels):\n",
        "        if self.cpu:\n",
        "            embeddings = embeddings.cpu()\n",
        "            \n",
        "        distance_matrix = pdist(embeddings)\n",
        "\n",
        "        labels = labels.cpu().data.numpy()\n",
        "        all_pairs = np.array(list(combinations(range(len(labels)), 2)))\n",
        "        all_pairs = torch.LongTensor(all_pairs)\n",
        "        positive_pairs = all_pairs[(labels[all_pairs[:, 0]] == labels[all_pairs[:, 1]]).nonzero()]\n",
        "        negative_pairs = all_pairs[(labels[all_pairs[:, 0]] != labels[all_pairs[:, 1]]).nonzero()]\n",
        "\n",
        "        negative_distances = distance_matrix[negative_pairs[:, 0], negative_pairs[:, 1]]\n",
        "        negative_distances = negative_distances.cpu().data.numpy()\n",
        "        top_negatives = np.argpartition(negative_distances, len(positive_pairs))[:len(positive_pairs)]\n",
        "        top_negative_pairs = negative_pairs[torch.LongTensor(top_negatives)]\n",
        "\n",
        "        return positive_pairs, top_negative_pairs"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "HXwFDac71E1I",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# 5.3 Online contrastive loss"
      ]
    },
    {
      "metadata": {
        "id": "dXreAeeb1E1J",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class OnlineContrastiveLoss(nn.Module):\n",
        "    \"\"\"\n",
        "    Online Contrastive loss\n",
        "    Takes a batch of embeddings and corresponding labels.\n",
        "    Pairs are generated using pair_selector object that take embeddings and targets and return indices of positive\n",
        "    and negative pairs\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, margin, pair_selector):\n",
        "        super(OnlineContrastiveLoss, self).__init__()\n",
        "        self.margin = margin\n",
        "        self.pair_selector = pair_selector\n",
        "\n",
        "    def forward(self, embeddings, target):\n",
        "        positive_pairs, negative_pairs = self.pair_selector.get_pairs(embeddings, target)\n",
        "        if embeddings.is_cuda:\n",
        "            positive_pairs = positive_pairs.cuda()\n",
        "            negative_pairs = negative_pairs.cuda()\n",
        "        positive_loss = (embeddings[positive_pairs[:, 0]] - embeddings[positive_pairs[:, 1]]).pow(2).sum(1)\n",
        "        negative_loss = F.relu(\n",
        "            self.margin - (embeddings[negative_pairs[:, 0]] - embeddings[negative_pairs[:, 1]]).pow(2).sum(\n",
        "                1).sqrt()).pow(2)\n",
        "        loss = torch.cat([positive_loss, negative_loss], dim=0)\n",
        "        return loss.mean()\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "FyF073ZH1E1L",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# 5.4 Training"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "goVi1_-PCuvL",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# We'll create mini batches by sampling labels that will be present in the mini batch and number of examples from each class\n",
        "\n",
        "train_batch_sampler = BalancedBatchSampler(train_dataset, num_classes=args.num_classes, num_samples=25)\n",
        "test_batch_sampler = BalancedBatchSampler(test_dataset, num_classes=args.num_classes, num_samples=25)\n",
        "\n",
        "kwargs = {'num_workers': args.num_workers, 'pin_memory': True} \n",
        "online_train_loader = torch.utils.data.DataLoader(train_dataset, batch_sampler=train_batch_sampler, **kwargs)\n",
        "online_test_loader = torch.utils.data.DataLoader(test_dataset, batch_sampler=test_batch_sampler, **kwargs)\n",
        "\n",
        "margin = 1.\n",
        "embedding_net = EmbeddingNet()\n",
        "model = embedding_net\n",
        "loss_fn = OnlineContrastiveLoss(margin, HardNegativePairSelector())\n",
        "\n",
        "model.to(args.device)\n",
        "loss_fn.to(args.device)\n",
        "\n",
        "args.lr = 1e-3\n",
        "optimizer = optim.Adam(model.parameters(), lr=args.lr)\n",
        "scheduler = lr_scheduler.StepLR(optimizer, 8, gamma=0.1, last_epoch=-1)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Jy5YAIBc1E1M",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def train_mining_epoch(train_loader, model, loss_fn, optimizer, args, log_interval=50):\n",
        "    model.train()\n",
        "    losses = []\n",
        "    total_loss, num_samples =  0, 0\n",
        "  \n",
        "    for batch_idx, (data, target) in enumerate(train_loader):\n",
        "        num_samples += data[0].size(0)\n",
        "        \n",
        "        if not type(data) in (tuple, list):\n",
        "            data = (data,)\n",
        "            \n",
        "        data = tuple(d.to(args.device) for d in data)\n",
        "        target = target.to(args.device)\n",
        "          \n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        outputs = model(*data)\n",
        "        \n",
        "        loss = loss_fn(outputs, target)\n",
        "        \n",
        "        losses.append(loss.data.item())\n",
        "\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        \n",
        "        if batch_idx % log_interval == 0:\n",
        "            print('Train: [{}/{} ({:.0f}%)]\\tLoss: {:.6f} '.format(\n",
        "                batch_idx * len(data[0]), len(train_loader.dataset),\n",
        "                100. * batch_idx / len(train_loader), np.mean(losses)))           \n",
        "            \n",
        "            total_loss += np.sum(losses)\n",
        "            losses = []\n",
        "            \n",
        "    return total_loss/(batch_idx + 1)\n",
        "\n",
        "def test_mining_epoch(test_loader, model, loss_fn, args, log_interval=50):\n",
        "    with torch.no_grad():\n",
        "        model.eval()\n",
        "        losses = []\n",
        "        num_samples = 0\n",
        "    \n",
        "        for batch_idx, (data, target) in enumerate(test_loader):\n",
        "\n",
        "            if not type(data) in (tuple, list):\n",
        "                data = (data,)\n",
        "                \n",
        "            data = tuple(d.to(args.device) for d in data)\n",
        "            target = target.to(args.device)\n",
        "            outputs = model(*data)\n",
        "\n",
        "            loss = loss_fn(outputs, target)\n",
        "            losses.append(loss.data.item())\n",
        "    \n",
        "        return np.sum(losses)/(batch_idx + 1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "MQnKNMAb1E1O",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "args.start_epoch = 0\n",
        "\n",
        "# needed for annealing learning rate in case of resuming of training\n",
        "for epoch in range(0, args.start_epoch):\n",
        "    scheduler.step()\n",
        "\n",
        "# main training loop\n",
        "for epoch in range(args.start_epoch, args.num_epochs):\n",
        "    scheduler.step()\n",
        "\n",
        "    # train stage\n",
        "    train_loss = train_mining_epoch(online_train_loader, model, loss_fn, optimizer, args)\n",
        "    message = 'Epoch: {}/{}. Train set: Average loss: {:.4f}'.format(\n",
        "        epoch + 1, args.num_epochs, train_loss)\n",
        "    \n",
        "    # testing/validation stage    \n",
        "    test_loss = test_mining_epoch(online_test_loader, model, loss_fn, args)\n",
        "    \n",
        "    message += '\\nEpoch: {}/{}. Validation set: Average loss: {:.4f}'.format(epoch + 1, args.num_epochs,\n",
        "                                                                             test_loss)\n",
        "    print(message)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "X13jYJfk1E1R",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## 5.5 Visualizations"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "ZCvF0AlCCuvX",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "train_embeddings_ocl, train_labels_ocl = extract_embeddings(train_loader, model, args)\n",
        "plot_embeddings(train_embeddings_ocl, train_labels_ocl, title='Train embeddings with hard negative mining')\n",
        "test_embeddings_ocl, test_labels_ocl = extract_embeddings(test_loader, model, args)\n",
        "plot_embeddings(test_embeddings_ocl, test_labels_ocl, title='Test embeddings with hard negative mining')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "QxPm_Zb81E1V",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# L2-normalize embeddings\n",
        "test_embeddings_ocl_norm = test_embeddings_ocl / np.linalg.norm(test_embeddings_ocl, axis=-1, keepdims=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "HXOI0bjf1E1W",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "test_images_raw = get_raw_images(test_loader)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "smvUHdXOyt7r",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "for i in range(5):\n",
        "    launch_query(test_embeddings_ocl_norm, test_images_raw)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}